from ultralytics import YOLO
import torch
from PIL import Image, ImageDraw
from torch import nn

if __name__ == '__main__':
    print(torch.cuda.device_count())
    # model = YOLO("yolov8n.pt")
    model = YOLO("runs/detect/train/weights/best.pt")
    model.export(format="onnx",imgsz=640,keras=True)
    # result = model.predict(source=Image.open("datasets/images/train/0.png"))
    # print("result")

    # print(result[0])
    # model.train(data="train.yaml", epochs=2,batch=4)  # train the model
    # metrics = model.val()  # evaluate model performance on the validation set
